# ECSB-2, "efficient sampling for combinatorial bandits". 
# Based on https://arxiv.org/abs/1502.03475. 

abstract type ECSB2OptimisationAlgorithm end
function optimise_linear_sqrtlinear(instance::CombinatorialInstance, ::ECSB2OptimisationAlgorithm, linear::Dict{T, Float64}, sqrtlinear::Dict{T, Float64}, epsilon::Float64, d::Int, verbose::Int) where T end

mutable struct ECSB2 <: Policy
  epsilon::Float64 # Only for approximation algorithm. 
  algo::ECSB2OptimisationAlgorithm
  verbose::Int # 0: nothing. 1: summary of function. 2: iteration counter too. 

  function ECSB2(epsilon::Float64, algo::ECSB2OptimisationAlgorithm, verbose::Int=0)
    return new(epsilon, algo, verbose)
  end
end

mutable struct ECSB2Details <: PolicyDetails
  nIterations::Int
  nBreakpoints::Int
  bestLambda::Float64
  bestLinearObjective::Float64
  bestNonlinearObjective::Float64
  solverTimes::Vector{Float64}

  function ECSB2Details()
    return new(0, 0, 0.0, 0.0, 0.0, Float64[])
  end
end

function choose_action(instance::CombinatorialInstance, policy::ECSB2, state::State{T}; with_trace::Bool=false) where T
  # For initialisation, if some arms have never been tried, force them to be tried (thus without linearity). 
  # Without this, the standard deviation term cannot exist. 
  if any(v == 0 for v in values(state.arm_counts))
    weights = Dict(arm => (state.arm_counts[arm] == 0.0) ? 1.0 : 0.0 for arm in keys(state.arm_counts))

    t0 = now()
    sol = solve_linear(instance, weights)
    t1 = now()

    if with_trace
      runDetails = ECSB2Details()
      push!(runDetails.solverTimes, (t1 - t0).value)
      return sol, runDetails
    else
      return sol
    end
  end

  # If all arms have been tried, can compute the two parts of the objective function: the average reward and a standard deviation. 
  arms = keys(state.arm_counts)
  t = state.arm_counts
  w = state.arm_average_reward
  n = state.round

  d = length(w)
  fn = log(n) + 4 * d * log(log(n))
  s2 = map(kv -> kv[1] => (fn / 2) / kv[2], t)

  # Solve the maximisation w^T x + \sqrt{s2^T x} using only the solve_linear oracle. 
  return optimise_linear_sqrtlinear(instance, policy.algo, w, s2, policy.epsilon, d, policy.verbose, with_trace=with_trace)
end
